package org.example.config;

import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.ollama.OllamaChatModel;
import dev.langchain4j.model.ollama.OllamaEmbeddingModel;
import dev.langchain4j.model.ollama.OllamaStreamingChatModel;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import org.example.entity.OllamaChatModelEntity;
import org.example.entity.PostgresSQLEntity;
import org.example.service.Assistant;
import org.example.service.ToolService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class AIConfig {

    @Autowired
    OllamaChatModelEntity ollamaChatModelEntity;

    @Autowired
    PostgresSQLEntity postgresSQLEntity;

    @Autowired
    ToolService toolService;

    @Bean
    public ChatLanguageModel chatLanguageModel() {
        return OllamaChatModel.builder()
                .baseUrl(ollamaChatModelEntity.getBaseUrl())
                .modelName(ollamaChatModelEntity.getModelName())
                .build();
    }

    @Bean
    public OllamaStreamingChatModel ollamaStreamingChatModel() {
        return OllamaStreamingChatModel.builder()
                .baseUrl(ollamaChatModelEntity.getBaseUrl())
                .modelName(ollamaChatModelEntity.getModelName())
                .build();
    }

    @Bean
    EmbeddingModel embeddingModel() {
        return OllamaEmbeddingModel.builder().baseUrl(ollamaChatModelEntity.getBaseUrl())
                .modelName(ollamaChatModelEntity.getModelName())
                .build();
    }

    @Bean
    ChatMemoryProvider chatMemoryProvider() {
        ChatMemoryProvider chatMemoryProvider = memoryId -> MessageWindowChatMemory.builder()
                .id(memoryId)
                .maxMessages(20)
                // 可以自己定义
                // .chatMemoryStore()
                .build();

        return chatMemoryProvider;
    }

    // 自定义向量仓库
    @Bean
    EmbeddingStore<TextSegment> embeddingStore() {
        // 基于内存的向量数据库
         return new InMemoryEmbeddingStore<>();
        // 基于 Postgres 数据库
//        return PgVectorEmbeddingStore.builder()
//                .createTable(true)
//                .dropTableFirst(true)
//                .database(postgresSQLEntity.getDatabase())
//                .host(postgresSQLEntity.getHost())
//                .port(postgresSQLEntity.getPort())
//                .user(postgresSQLEntity.getUser())
//                .password(postgresSQLEntity.getPassword())
//                .table(postgresSQLEntity.getTable())
//                .dimension(384)
//                .build();
    }

    @Bean
    Assistant assistant() {
        return AiServices.builder(Assistant.class)
                // 设置 LLM 模型
                .chatLanguageModel(chatLanguageModel())
                // 设置流传输的 LLM 模型
                .streamingChatLanguageModel(ollamaStreamingChatModel())
                // 会话记忆，最多为 20 条
                .chatMemoryProvider(chatMemoryProvider())
                // 向量数据库
                .contentRetriever(EmbeddingStoreContentRetriever.from(embeddingStore()))
                // deepseek 不支持函数回调
                // .tools(toolService)
                .build();
    }

}
